from helpFunc import *
from pyNetns import *
import select, time
import socks,socket
from collections import deque
import uuid,threading
import logging
from logging.handlers import RotatingFileHandler
from datetime import datetime

logger = logging.getLogger('a')
handler = RotatingFileHandler('a.log', maxBytes=15 * 1024, backupCount=1)
formatter = logging.Formatter("%(asctime)s %(message)s",
                              "%Y-%m-%d %H:%M:%S")
handler.setFormatter(formatter)
logger.addHandler(handler)

cutThrOn = False
cmdFile = 'a_cmd.txt'
os.system('echo start>%s'%cmdFile)

class main():
    def __init__(self):
        self.tap0 = openTap(b'tap0')
        self.tap1 = openTap(b'tap1')
        self.tap2 = openTap(b'tap2')

        self.udpListenSock = udpListenSock = nsUdpSock('a')
        udpListenSock.setsockopt(socket.SOL_IP, 19, 1)
        udpListenSock.setsockopt(socket.IPPROTO_IP, 20, 1)
        udpListenSock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        udpListenSock.bind(('0.0.0.0', 8081))

        self.tcpListenSock = tcpListenSock = nsTcpSock('a')
        tcpListenSock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        tcpListenSock.bind(('0.0.0.0', 8080))
        tcpListenSock.listen(500)

        self.macMap = {}
        self.fackIpMap = {}
        self.ipCo = 0
        self.ipCoQ = deque()

        self.udpAddMap = {}
        self.udpSerSockMap = {}
        self.tcpIdMap = {}
        self.tcpSerSockMap = {}
        self.tcpCliSockMap = {}
        self.connectPollingMap = {}

        self.lastCutIpTime = 0
        self.dataSend = self.dataRecv = 0
        self.staTime = 0
        self.cutThread = None
        self.ipInCut = []


    def dealUdpBack(self, fd, currentTime):
        k = self.udpSerSockMap[fd]
        v = self.udpAddMap[k]
        srcAdd, dstAdd = k
        isDebug = debugIp in dstAdd[0]
        s = v['serSock']
        sock = v['cliSock']
        try:
            d, a = s.recvfrom(sockReadLen)
            sock.sendto(d, srcAdd)
            self.dataRecv += len(d)
            if isDebug:
                debugPrint('UdpBack dataLen', k, len(d))
        except Exception as e:
            if isDebug:
                debugPrint('UdpBack error', k)
            print(time.monotonic(), 'excp dealUdpBack', e)
            return
        self.udpAddMap[k]['time'] = currentTime

    def dealUdp(self, currentTime):
        msg, ancdata, flags, srcAdd = self.udpListenSock.recvmsg(100000, 100)
        for data in ancdata:
            if data[0] == 0 and data[1] == 20:
                port = data[2][2] * 256 + data[2][3]
                ipB = data[2][4:8]
                ip = socket.inet_ntoa(ipB)
                dstAdd = (ip, port)
                isDebug = debugIp in ip

                if (srcAdd, dstAdd) not in self.udpAddMap:
                    if isDebug:
                        debugPrint('dealUdp and make socket', (srcAdd, dstAdd))
                    if enableUdpProxy:
                        s = socks.socksocket(type=socket.SOCK_DGRAM)
                        s.set_proxy(socks.SOCKS5, sslocalAddUdp[0], sslocalAddUdp[1])
                    else:
                        s = socket.socket(type=socket.SOCK_DGRAM)
                    sock = nsUdpSock('a')
                    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                    sock.bind(dstAdd)
                    self.udpAddMap[(srcAdd, dstAdd)] = {'serSock': s, 'cliSock': sock, 'time': currentTime}
                    self.udpSerSockMap[s.fileno()] = (srcAdd, dstAdd)
                else:
                    if isDebug:
                        debugPrint('dealUdp and use old socket', (srcAdd, dstAdd))
                s = self.udpAddMap[(srcAdd, dstAdd)]['serSock']
                try:
                    s.sendto(msg, dstAdd)
                    self.dataSend += len(msg)
                    if isDebug:
                        debugPrint('dealUdp and send ok', len(msg))
                except Exception as e:
                    if isDebug:
                        debugPrint('dealUdp and send err', len(msg))
                    print(time.monotonic(), 'excp dealUdp', e)
                    return
                self.udpAddMap[(srcAdd, dstAdd)]['time'] = currentTime

    def startCutThr(self,l):
        global cutThrOn
        cutThrOn = True        
        def ff(l):
            global cutThrOn
            for si in l: 
                #os.system('ip netns exec a ifconfig>a_temp2')
                #f = open('a_temp2')
                #ss = f.read()
                #f.close()
                #if 'mydev:%s:'%si not in ss:
                    #print('error',si)
                cmd = 'ip netns exec a ifconfig mydev:%s down >a_temp 2>&1' % (si,)
                os.system('echo "%s">>%s'%(cmd,cmdFile))
                os.system(cmd)
                os.system('cat a_temp>>%s'%(cmdFile,))
                #f = open('a_temp')
                #ss = f.read()
                #f.close()
                #if ss:
                    #print('error',cmd)
            doIproute()
            print(datetime.now(),'cutThread over')            
            cutThrOn = False
            
        t = threading.Thread(target=ff,args=(l,))
        t.setDaemon(True)
        t.start()
        self.cutThread = t                            
    
    def cupFakeIp(self, currentTime):    
        print (len(self.fackIpMap),'fackIpMap len')    
        if cutThrOn or len(self.fackIpMap) < maxFakeIp or currentTime - self.lastCutIpTime < cupIpInterval:
            return
        cutL = []
        self.ipInCut = []
        self.lastCutIpTime = currentTime
        print(datetime.now(), 'cut before', len(self.fackIpMap))

        for dstip in list(self.fackIpMap.keys()):
            dstipStr = socket.inet_ntoa(dstip)
            isDebug = debugIp in dstipStr
            if time.monotonic() - currentTime > 1:
                print(datetime.now(), 'cut after', len(self.fackIpMap))   
                self.startCutThr(cutL)
                return

            if currentTime - self.fackIpMap[dstip]['time'] < fakeIpEffectTime:
                if isDebug:
                    debugPrint(debugIp, 'cupFakeIp time new')
                continue

            found = False
            for k, v in self.udpAddMap.items():
                if dstipStr in k[1][0]:
                    found = True
                    if isDebug:
                        debugPrint(debugIp, 'cupFakeIp udp found')
                    break
            if found:
                continue

            for k, v in self.connectPollingMap.items():
                if dstipStr in v['dstAdd'][0]:
                    found = True
                    if isDebug:
                        debugPrint(debugIp, 'cupFakeIp poll found')
                    break
            if found:
                continue

            for k, v in self.tcpIdMap.items():
                if dstipStr in v['dstAdd'][0]:
                    found = True
                    if isDebug:
                        debugPrint(debugIp, 'cupFakeIp tcp found')
                    break
            if found:
                continue

            si = self.fackIpMap[dstip]['ipCo']
            self.ipInCut.append(dstip)
            cutL.append(si)
            self.ipCoQ.append(si)
            del self.fackIpMap[dstip]
            if isDebug:
                debugPrint(debugIp, 'cupFakeIp cut', si)
        print(datetime.now(), 'cut after', len(self.fackIpMap))
        self.startCutThr(cutL)

    def refreshMac(self, st, currentTime):
        desMac = st[:6]
        srcMac = st[6:12]
        srcIpB = st[26:30]
        dstIpB = st[30:34]
        srcIp = socket.inet_ntoa(srcIpB)
        desIp = socket.inet_ntoa(dstIpB)
        isDebug = debugIp in desIp
        
        if cutThrOn and dstIpB in self.ipInCut:
            print('drop ipInCut',desIp)
            return False

        if mydevIp3 not in srcIp:
            return False

        if mydevIp3 in desIp:
            return False

        if desMac != con_devMacB or desIp == con_gatewayIp or srcMac == b'\xff' * 6 \
                or dstIpB[0] >= 224:
            return False

        if self.macMap.get(srcIpB, b'') != srcMac:
            self.macMap[srcIpB] = srcMac

        if dstIpB not in self.fackIpMap:
            if self.ipCoQ and not cutThrOn:
                si = self.ipCoQ.popleft()
            else:
                self.ipCo += 1
                si = self.ipCo

            self.fackIpMap[dstIpB] = {'time': currentTime, 'ipCo': si}
            cmd = "ip netns exec a ifconfig mydev:%s %s/0 up >>%s 2>&1" % (si, desIp,cmdFile)
            os.system('echo "%s">>%s'%(cmd,cmdFile))
            os.system(cmd)
            doIproute()
            if isDebug:
                debugPrint('add macMap', si)
        else:
            self.fackIpMap[dstIpB]['time'] = currentTime

        return True

    def dealTap1(self, currentTime):
        st = os.read(self.tap1.fileno(), tapReadLen)
        ty, srcIp, desIp, flags, code = checkEth(st)

        if not ty:
            return
        isDebug = debugIpB in srcIp
        if isDebug:
            debugPrint('tap1', ty, socket.inet_ntoa(srcIp), socket.inet_ntoa(desIp), flags, code)

        if desIp in self.macMap and (ty == ty_tcp or ty == ty_udp):
            if srcIp in self.fackIpMap:
                self.fackIpMap[srcIp]['time'] = currentTime
            os.write(self.tap0.fileno(), self.macMap[desIp] + st[6:])

    def dealTap2(self, currentTime):
        st = os.read(self.tap2.fileno(), tapReadLen)
        ty, srcIp, dstIp, flags, code = checkEth(st)
        if not ty:
            return

        if srcIp != mydevIpB:
            return

        if ty == ty_arp_request or ty == ty_arp_reply or ty == ty_icmp:
            os.write(self.tap0.fileno(), st)

    def dealTap0(self, currentTime):
        st = os.read(self.tap0.fileno(), tapReadLen)
        ty, srcIp, dstIp, flags, code = checkEth(st)
        if not ty:
            return
        isDebug = debugIpB in dstIp
        if isDebug:
            debugPrint('tap0', ty, socket.inet_ntoa(srcIp), socket.inet_ntoa(dstIp), flags, code)
        if (ty == ty_arp_request or ty == ty_arp_reply or ty == ty_icmp) and dstIp == mydevIpB:
            os.write(self.tap2.fileno(), st)

        elif ty == ty_tcp or ty == ty_udp:
            if self.refreshMac(st, currentTime):
                os.write(self.tap1.fileno(), con_devMacB + con_fakeMacB + st[12:])

    def dealTcp(self):
        csock, srcAdd = self.tcpListenSock.accept()
        csock.setblocking(0)
        orig_dst = csock.getsockopt(socket.SOL_IP, 80, 16)
        des_port = struct.unpack('>H', orig_dst[2:4])[0]
        des_ip = socket.inet_ntoa(orig_dst[4:8])
        isDebug = debugIp in des_ip

        dstAdd = (des_ip, des_port)
        if isDebug:
            debugPrint('tcp conn', srcAdd, dstAdd)

        s = socks.socksocket()
        s.set_proxy(socks.SOCKS5, sslocalAdd[0], sslocalAdd[1])
        s.setblocking(0)
        s.connect_ex(dstAdd)
        self.connectPollingMap[s.fileno()] = {'cliSock': csock, 'serSock': s, 'srcAdd': srcAdd, 'dstAdd': dstAdd}

    def closeTcp(self, id):
        s1 = self.tcpIdMap[id]['serSock']
        s2 = self.tcpIdMap[id]['cliSock']
        del self.tcpIdMap[id]
        del self.tcpSerSockMap[s1.fileno()]
        del self.tcpCliSockMap[s2.fileno()]
        s1.close()
        s2.close()

    def rigTcp(self, po):
        for id in list(self.tcpIdMap.keys()):
            v = self.tcpIdMap[id]
            isDebug = debugIp in v['dstAdd'][0]
            if v['cliErr']:
                if isDebug:
                    debugPrint('tcp close cliErr')
                self.closeTcp(id)
                continue
            if v['serErr'] and not v['serData']:
                if isDebug:
                    debugPrint('tcp close serErr and no data')
                self.closeTcp(id)
                continue
            if v['serErr']:
                po.register(v['cliSock'], select.POLLOUT)
                continue

            if len(v['cliData']) < con_cache:
                cliIn = select.POLLIN
            else:
                cliIn = 0

            if len(v['serData']) < con_cache:
                serIn = select.POLLIN
            else:
                serIn = 0

            if v['serData']:
                cliOut = select.POLLOUT
            else:
                cliOut = 0

            if v['cliData']:
                serOut = select.POLLOUT
            else:
                serOut = 0
            po.register(v['cliSock'], cliIn | cliOut)
            po.register(v['serSock'], serIn | serOut)

    def dealCli(self, fd, e, currentTime):
        id = self.tcpCliSockMap[fd]
        isDebug = debugIp in self.tcpIdMap[id]['dstAdd']

        if e & select.POLLERR:
            if isDebug:
                debugPrint('cliSock poll error')
            self.tcpIdMap[id]['cliErr'] = True
            return

        if e & select.POLLIN:
            try:
                ss = self.tcpIdMap[id]['cliSock'].recv(sockReadLen)
                if not ss:
                    if isDebug:
                        debugPrint('cliSock data in err')
                    self.tcpIdMap[id]['cliErr'] = True
                    return
                self.tcpIdMap[id]['cliData'] += ss
                self.tcpIdMap[id]['time'] = currentTime
                if isDebug:
                    debugPrint('cliSock data in len', len(ss))
            except:
                if isDebug:
                    debugPrint('cliSock data in err')
                self.tcpIdMap[id]['cliErr'] = True

        if e & select.POLLOUT:
            try:
                n = self.tcpIdMap[id]['cliSock'].send(self.tcpIdMap[id]['serData'])
                self.tcpIdMap[id]['serData'] = self.tcpIdMap[id]['serData'][n:]
                self.tcpIdMap[id]['time'] = currentTime
                if isDebug:
                    debugPrint('cliSock data out len', n)
            except:
                self.tcpIdMap[id]['cliErr'] = True
                if isDebug:
                    debugPrint('cliSock data out err')

    def dealSer(self, fd, e, currentTime):
        id = self.tcpSerSockMap[fd]
        isDebug = debugIp in self.tcpIdMap[id]['dstAdd']

        if e & select.POLLERR:
            if isDebug:
                debugPrint('serSock poll err')
            self.tcpIdMap[id]['serErr'] = True
            return

        if e & select.POLLIN:
            try:
                ss = self.tcpIdMap[id]['serSock'].recv(sockReadLen)
                if not ss:
                    self.tcpIdMap[id]['serErr'] = True
                    if isDebug:
                        debugPrint('serSock data in err')
                    return
                self.tcpIdMap[id]['serData'] += ss
                self.tcpIdMap[id]['time'] = currentTime
                self.dataRecv += len(ss)
                if isDebug:
                    debugPrint('serSock data in len', len(ss))
            except:
                if isDebug:
                    debugPrint('serSock data in err')
                self.tcpIdMap[id]['serErr'] = True

        if e & select.POLLOUT:
            try:
                n = self.tcpIdMap[id]['serSock'].send(self.tcpIdMap[id]['cliData'])
                self.tcpIdMap[id]['cliData'] = self.tcpIdMap[id]['cliData'][n:]
                self.tcpIdMap[id]['time'] = currentTime
                self.dataSend += n
                if isDebug:
                    debugPrint('serSock data out len', n)
            except:
                if isDebug:
                    debugPrint('serSock data out err')
                self.tcpIdMap[id]['serErr'] = True

    def cutTcpMap(self, currentTime):
        l = list(self.tcpIdMap.keys())
        for id in l:
            if currentTime - self.tcpIdMap[id]['time'] > tcpEffctTime:
                if debugIp in self.tcpIdMap[id]['dstAdd'][0]:
                    debugPrint('tcp close inactive time')

                self.closeTcp(id)

    def cutUdpMap(self, currentTime):
        l = list(self.udpAddMap.keys())
        for i in l:
            if currentTime - self.udpAddMap[i]['time'] > udpEffctTime:
                if debugIp in i[1][0]:
                    debugPrint('udp close inactive time')
                del self.udpSerSockMap[self.udpAddMap[i]['serSock'].fileno()]
                self.udpAddMap[i]['serSock'].close()
                self.udpAddMap[i]['cliSock'].close()
                del self.udpAddMap[i]

    def periodWork(self, currentTime):
        self.cutTcpMap(currentTime)
        self.cutUdpMap(currentTime)
        self.cupFakeIp(currentTime)
        st = 'tcp %s udp %s poll %s send %sKB recv %sKB' \
             % (len(self.tcpIdMap), len(self.udpAddMap), len(self.connectPollingMap),
                self.dataSend // 1024, self.dataRecv // 1024)
        # print(currentTime,st)
        logger.warning(st)
        self.dataSend = self.dataRecv = 0
        doIproute_noblock()

    def dealConnPolling(self, fd, e):
        s = self.connectPollingMap[fd]['cliSock']
        serS = self.connectPollingMap[fd]['serSock']
        srcAdd = self.connectPollingMap[fd]['srcAdd']
        dstAdd = self.connectPollingMap[fd]['dstAdd']
        isDebug = debugIp in dstAdd[0]
        if e & select.POLLERR:
            if isDebug:
                debugPrint('conn error')
            del self.connectPollingMap[fd]
            serS.close()
            s.close()
        else:
            if isDebug:
                debugPrint('conn ok')
            id = str(uuid.uuid1())
            self.tcpIdMap[id] = {'serSock': serS, 'cliSock': s, 'time': time.monotonic(),
                                 'serData': b'', 'cliData': b'', 'serErr': False, 'cliErr': False,
                                 'srcAdd': srcAdd, 'dstAdd': dstAdd}
            self.tcpSerSockMap[serS.fileno()] = id
            self.tcpCliSockMap[s.fileno()] = id
            del self.connectPollingMap[fd]

    def work(self):
        currentTime = time.monotonic()
        if currentTime - self.staTime > 1:
            self.staTime = currentTime
            self.periodWork(currentTime)

        po = select.poll()
        po.register(self.tap0, select.POLLIN)
        po.register(self.tap1, select.POLLIN)
        po.register(self.tap2, select.POLLIN)
        po.register(self.udpListenSock, select.POLLIN)
        po.register(self.tcpListenSock, select.POLLIN)

        for one in self.udpSerSockMap:
            po.register(one, select.POLLIN)

        for one in self.connectPollingMap:
            po.register(one, select.POLLOUT)

        self.rigTcp(po)
        es = po.poll(1000)

        for fd, e in es:
            if fd in self.connectPollingMap:
                self.dealConnPolling(fd, e)

            elif fd in self.tcpCliSockMap:
                self.dealCli(fd, e, currentTime)

            elif fd in self.tcpSerSockMap:
                self.dealSer(fd, e, currentTime)

            elif fd in self.udpSerSockMap:
                self.dealUdpBack(fd, currentTime)

            elif fd == self.tcpListenSock.fileno():
                self.dealTcp()

            elif fd == self.udpListenSock.fileno():
                self.dealUdp(currentTime)

            elif fd == self.tap0.fileno():
                self.dealTap0(currentTime)

            elif fd == self.tap1.fileno():
                self.dealTap1(currentTime)

            elif fd == self.tap2.fileno():
                self.dealTap2(currentTime)

if __name__ == "__main__":
    ini()
    m = main()
    print('start')
    while True:
        m.work()
